# Copyright 2012 Google Inc. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Integration test for gaussian_contributors with variational inference."""

__author__ = 'tpw@google.com (Tamsyn Waterhouse)'

import unittest
from resolver import alternating_resolution
from resolver import gaussian_contributors
from resolver import test_util

# The expected variance is the same for every question, because every question
# was answered by the same multiset of contributors; this is a property of the
# model.
EXPECTED_VARIANCE = 0.019

# These values were generated by iterating until convergence, and they
# correspond well to the results from using a confusion matrix model (and also
# to the results from running EM over the Gaussian model):
EXPECTED_MEAN = {1: 1.000, 2: 3.209, 3: 1.508, 4: 1.861, 5: 2.093,
                 6: 2.210, 7: 1.468, 8: 3.117, 9: 2.115, 10: 2.303,
                 11: 4.000, 12: 2.551, 13: 1.000, 14: 1.981, 15: 1.187,
                 16: 1.093, 17: 1.000, 18: 1.000, 19: 1.885, 20: 2.024,
                 21: 2.000, 22: 1.885, 23: 2.093, 24: 1.813, 25: 1.000,
                 26: 1.000, 27: 2.187, 28: 1.000, 29: 1.000, 30: 1.300,
                 31: 1.000, 32: 2.883, 33: 1.000, 34: 2.000, 35: 2.206,
                 36: 3.393, 37: 1.925, 38: 2.508, 39: 3.002, 40: 1.000,
                 41: 1.000, 42: 1.280, 43: 2.187, 44: 1.187, 45: 2.000}


class EmConfusionTest(unittest.TestCase):
  longMessage = True

  def testIterateUntilConvergence(self):
    # Initialize a Gaussian model and an EM object:
    gc = gaussian_contributors.GaussianContributors()
    vb = alternating_resolution.VariationalBayes()

    # Use the judgments from the Dawid & Skene example:
    data = test_util.DS_DATA
    gc.InitializeResolutions(data)
    # The algorithm converges slowly because EPSILON is small:
    self.assertFalse(vb.IterateUntilConvergence(data, gc, max_iterations=1000))
    expected = {}
    for question in data:
      expected[question] = dict([
          (gaussian_contributors.MEAN, EXPECTED_MEAN[question]),
          (gaussian_contributors.VARIANCE, EXPECTED_VARIANCE)])
    result = gc.ExtractResolutions(data)
    test_util.AssertResolutionsAlmostEqual(self, expected, result)


if __name__ == '__main__':
  unittest.main()
